#!/usr/bin/env python
# Script to read in profiles from bufkit in Profile format
# and plot the forecast times
import os, re, sys
from datetime import datetime,timedelta
import pickle
from optparse import OptionParser
import numpy as np

from INFILE import siteid, webdir, UTC_offset, maindir
sys.path.append('%s/data_parsers' % maindir)

# Get the station ID we want from the command line
parser = OptionParser()
parser.add_option('-s', '--site', dest='siteid', action='store', type='string', default=siteid, help='4-letter id of desired site')
parser.add_option('-e', '--export', dest='export', action='store_true',default=False, help='Write images to png or not.')
(opts,args) = parser.parse_args()
siteid = opts.siteid.lower()
export_flag = opts.export


print "Importing surface_parse_bufkit"
from surface_parse_bufkit import bufkit_parser, obs_parser, mos_parser, lamp_parser
print "Done"

print "Importing matplotlib"
import matplotlib
if export_flag:
    matplotlib.use('agg')

print "Importing pylab"
import pylab

#siteid = 'keri'
# outdir is the directory to move the files to if exporting
#outdir = '/home/disk/pvort/lmadaus/public_html/bufkit/plots'
outdir = webdir


# Only interested in the 24-hour period we're forecasting for
# with a lead-in time.  ftime is the actual times we're
# forecasting between, ptime is the times we're plotting

nowtime = datetime.now()
stime = nowtime.replace(hour=6,minute=0,second=0,microsecond=0)
stime = stime + timedelta(hours=24)
print "HERE", stime.strftime('%Y%m%d%H')

start_ftime = stime
# this must change depending on UTC offset
start_ftime_wind = stime + timedelta(hours=(UTC_offset-6))

lead_time = timedelta(hours=18)
fcst_window = timedelta(hours=24)
end_ftime = start_ftime + fcst_window
end_ftime_wind = start_ftime_wind + timedelta(hours=24)

start_ptime = start_ftime - lead_time
end_ptime = end_ftime + timedelta(hours=6)

# Now make a list of times to plot
timelist = []
t = start_ptime
fcst_delta = timedelta(hours=1)
while t <= end_ptime:
    timelist.append(t)
    t = t + fcst_delta

# Grab the most recent files
# Use Penn State repository for profiles
# Clear all old directories
os.system('rm *.buf* *.buz*')
#NAM_addr = 'ftp://ftp.meteo.psu.edu/pub/bufkit/nam_%s.buf' % siteid
#NAMM_addr = 'ftp://ftp.meteo.psu.edu/pub/bufkit/namm_%s.buf' % siteid
NAM_addr = 'http://www.meteor.iastate.edu/~ckarsten/bufkit/data/nam/nam_%s.buf' % siteid
NAMM_addr = 'http://www.meteor.iastate.edu/~ckarsten/bufkit/data/namm/namm_%s.buf' % siteid
GFS00Z_addr = 'ftp://ftp.meteo.psu.edu/pub/bufkit/GFS/00/gfs3_%s.buf' % siteid
GFS06Z_addr = 'ftp://ftp.meteo.psu.edu/pub/bufkit/GFS/06/gfs3_%s.buf' % siteid
GFS12Z_addr = 'ftp://ftp.meteo.psu.edu/pub/bufkit/GFS/12/gfs3_%s.buf' % siteid
GFS18Z_addr = 'ftp://ftp.meteo.psu.edu/pub/bufkit/GFS/18/gfs3_%s.buf' % siteid
RAP15Z_addr = 'ftp://ftp.meteo.psu.edu/pub/bufkit/RAP/15/rap_%s.buf' % siteid
RAP18Z_addr = 'ftp://ftp.meteo.psu.edu/pub/bufkit/RAP/18/rap_%s.buf' % siteid
#EWRF00Z_addr = 'http://wrfensemble.wdtb.noaa.gov/00/ewrf_%s.buz' % siteid
#EWRF12Z_addr = 'http://wrfensemble.wdtb.noaa.gov/12/ewrf_%s.buz' % siteid

filelist = []
types = ['NAM','GFS','RAP','EWRF']
# First, append nam files
filelist.append(NAM_addr)
filelist.append(NAMM_addr)
# Then GFS files
filelist.append(GFS18Z_addr)
filelist.append(GFS12Z_addr)
filelist.append(GFS06Z_addr)
filelist.append(GFS00Z_addr)
# Then RAP files
filelist.append(RAP18Z_addr)
filelist.append(RAP15Z_addr)

# Then EWRF means
#filelist.append(EWRF12Z_addr)
#filelist.append(EWRF00Z_addr)

model_ids = {'nam'      : '12Z NAM',
             'namm'     : '18Z NAM',
             'namm_p1'  : '06Z NAM',
             'gfs3'     : '18Z GFS',
             'gfs3_p1'  : '12Z GFS',
             'gfs3_p2'  : '06Z GFS',
             'gfs3_p3'  : '00Z GFS',
             'rap'      : '18Z RAP',
             'nam_mos'  : 'nam_mos',
             'gfs_mos'  : 'gfs_mos',
             'rap_p1'   : '15Z RAP',
             'usl'      : '12Z USL',
             'usl_p1'   : '22Z USL'}




for file in filelist:
     os.system('wget %s -nv' % file)
     #if '.buz' in file:
     #    filename = re.search('/(\D{4}_\D{4}.buz.?\d?)',file).groups()[0]
     #    os.system('mv %s %s.gz'% (filename,filename))
     #    os.system('gunzip %s.gz' % filename)
     #    #newfile = file.replace('.buz','.buf')
     #    #os.system('mv %s %s' % (file,newfile))


# Get a list of all files in the directory
# Current directory, just to make sure all downloaded
# correctly
cwd = os.getcwd()
allfiles = os.listdir(cwd)
# Only get the bufkit files
filelist = [f for f in allfiles if '.buf' in f]
print filelist

# We assume we've grabbed the most recent files
# from each model and there is some overlap.  There
# should only be one profile for each model.
# Filenames should start with the model they are 
# representing. Get a dictionary of model names.
models = {}
# Import the timeheight function
from plot_bufkit_timeheight import plot_timeheight

for file in filelist:
    # Use a regular expression to grab the model and site
    linesplit = re.search('(.*)_(.*)(.buf|.buz).?(\d?)',file).groups()
    modelid = linesplit[0]
    print "MODEL:", modelid
    print linesplit
    if linesplit[3] != '':
      modelid = modelid + '_p%s' % linesplit[3]
    # Call the model parser.  Its ouput will be
    # a dictionary of model forecast times.
    profile = bufkit_parser(file)
    models[modelid] = profile
    # Plot the time-heights appropriately
    if modelid in ['nam','namm','gfs3','gfs3_p1','gfs3_p2','gfs3_p3']:
        prev_modelid = modelid
        if modelid.startswith('gfs3'):
            modelid = 'gfs'
        print "Plotting timeheight for:", modelid
        plot_timeheight(file,'SKNT',modelid,siteid.upper(),outdir,'%s/site_data'%maindir)
        plot_timeheight(file,'TMPC',modelid,siteid.upper(),outdir, '%s/site_data'%maindir)
        plot_timeheight(file,'DWP-DEP',modelid,siteid.upper(),outdir, '%s/site_data'%maindir)
        plot_timeheight(file,'OMEG',modelid,siteid.upper(),outdir, '%s/site_data'%maindir)
        if not modelid.startswith('gfs'):
            plot_timeheight(file,'CFRL',modelid.upper(),siteid.upper(),outdir, '%s/site_data'%maindir)
        modelid = prev_modelid


#print models.keys()
#print models['nam'].keys()


# Added here to plot the BL winds which come from the timeheight
from plot_bufkit_blwinds import plot_current_blwinds
plot_current_blwinds(siteid,outdir,UTC_offset,True)

# Slim down the dictionaries to the times we need
for model in models.keys():
    for t in models[model].keys():
        if t not in timelist:
            del models[model][t]
#print keylist

def C_to_F(TC):
    # Convert Celsius to Fahrenheit
    TF = TC * 9. / 5. + 32
    return TF
def knt_to_mph(SK):
    # Convert Knots to MPH
    SM = SK * 1.15077945
    return SM
def mm_to_in(RM):
    # Convert millimeters to inches
    RI = RM * 0.03937
    return RI


# Make a dictionary for plotting parameters
color_dict = {'gfs3'    : 'midnightblue',
              'gfs3_p1' : 'blue',
              'gfs3_p2' : 'cornflowerblue',
              'gfs3_p3' : 'cyan',
              'nam'     : 'limegreen',
              'namm'    : 'darkgreen',
              'namm_p1' : 'lawngreen',
              'rap'     : 'firebrick',
              'rap_p1'  : 'red',
              'gfs_mos' : 'blue',
              'nam_mos' : 'green',
              'usl'     : 'gold',
              'usl_p1'  : 'GoldenRod',
              'GFS LAMP': 'magenta'}


def plot_wind_barbs(modlist,moslist):
    pylab.figure(figsize=(12,11))
    matplotlib.rcParams['xtick.minor.pad']=15
    # Loop through the models and get the U and V winds
    # First sort the models by time and name
    modids = {}
    for model in modlist:
        if model not in ['GFS LAMP']:
            moddates = models[model].keys()
            moddates.sort()
            startdate = moddates[0]
            if model == 'gfs3_p2' and startdate < start_ptime+timedelta(hours=2):
                startdate = moddates[0] - timedelta(hours=6)
            if model == 'gfs3' and startdate < start_ptime+timedelta(hours=2):
                startdate = start_ptime - timedelta(hours=18) 
            elif model == 'gfs3_p3' and startdate < start_ptime+timedelta(hours=2):
                startdate = moddates[0] - timedelta(hours=12)
            elif model == 'usl':
                startdate = moddates[0] - timedelta(hours=18)
            elif model == 'usl_p1':
                startdate = moddates[0] - timedelta(hours=8)
            elif model == 'namm' and startdate < start_ptime + timedelta(hours=2):
                startdate = start_ptime - timedelta(hours=6)
                model = 'namm_p1'
            if model.startswith('rap') and moddates[0].hour < 14:
                pass
            else:
                modids['%s_%s' % (startdate.strftime('%Y%m%d%H'),model)] = model
    # Now loop through each model in the sorted keys from newest to oldest
    modsorted = modids.keys()
    modsorted.sort()
    modsorted.reverse()
    keynums = range(len(modsorted))
    keynums.reverse()
    for modid,keynum in zip(modsorted,keynums):
        model = modids[modid]
        # Get a sorted list of dates
        if model == 'namm_p1':
            moddates = models['namm'].keys()
        else:
            moddates = models[model].keys()
        moddates.sort()
        # Get the u and v winds
        if model == 'namm_p1':
            uwnd = [models['namm'][d].uwnd for d in moddates]
            vwnd = [models['namm'][d].vwnd for d in moddates]
        else:
            uwnd = [models[model][d].uwnd for d in moddates]
            vwnd = [models[model][d].vwnd for d in moddates]
        # Plot the barbs on the line
        plotmoddates = [matplotlib.dates.date2num(d) for d in moddates]
        pylab.barbs(plotmoddates,[keynum]*len(moddates),uwnd,vwnd,barbcolor=color_dict[model],flagcolor=color_dict[model])
        pylab.hold(True)

    # Some things to make the figure look nice
    pylab.title('Forecast Wind Direction at %s' % (siteid.upper()))
    pylab.xlabel('Valid Time')
    pylab.ylabel('Model')
    # Make the yticks the model names
    keynums.reverse()
    modsorted.reverse()
    pylab.yticks(keynums,[model_ids[modids[m]] for m in modsorted])

    # Grey out non-forecast period
    pylab.axvspan(matplotlib.dates.date2num(start_ptime),matplotlib.dates.date2num(start_ftime),facecolor='0.8',alpha=0.40)
    pylab.axvspan(matplotlib.dates.date2num(end_ftime),matplotlib.dates.date2num(end_ptime),facecolor='0.8',alpha=0.40)

    ax = pylab.gca()
    # Reformat the date axis
    ax.set_xlim([matplotlib.dates.date2num(start_ptime),matplotlib.dates.date2num(end_ptime)])
    #min,max = ax.get_ylim()

    ax.xaxis.set_major_locator(matplotlib.dates.HourLocator(byhour=[0,3,6,9,12,15,18,21]))
    ax.xaxis.set_major_formatter(matplotlib.dates.DateFormatter('%HZ'))
    ax.xaxis.set_minor_locator(matplotlib.dates.DayLocator())
    ax.xaxis.set_minor_formatter(matplotlib.dates.DateFormatter('%h %d'))
    pylab.grid()
    if export_flag:
        pylab.savefig('%s_winddir.png' % (siteid.upper()),bbox_inches='tight')
        os.system('mv %s_winddir.png %s/%s_winddir.png' % (siteid.upper(),outdir,siteid.upper()))
    else:
        pylab.show()
 
def plot_models(var,modlist,oblist,moslist):
    # Now we can get down to plotting.  First, plot the temperature
    pylab.figure(figsize=(12,11))
    matplotlib.rcParams['xtick.minor.pad']=15
    # Accumulate sorted lists
    max_val = None
    min_val = None
    precip_fcst = {}
    wind_fcsts = {}
    for model in modlist+moslist:
        var_list = []
        # get a sorted list of keys
        keysorted = models[model].keys()
        keysorted.sort()
        # Now get the variable for each key
        for date in keysorted:
            if var == 'temp':
                var_list.append(C_to_F(models[model][date].tmpc))
            elif var == 'wspd':
                var_list.append(models[model][date].sknt)
                try:
                    wind_fcsts[model][date] = models[model][date].sknt
                except:
                    wind_fcsts[model] = {}
                    wind_fcsts[model][date] = models[model][date].sknt
            elif var == 'precip':
                if date <= start_ftime:
                    var_list.append(0.0)
                else:
                    if model.endswith('mos'):
                        var_list.append(0.0)
                    else:
                        if model == 'usl':
                            #USL is already in inches
                            var_list.append(var_list[-1]+models[model][date].p01m)
                        else:
                            var_list.append(var_list[-1]+mm_to_in(models[model][date].p01m))
                # This is the one place we can get the model precip forecasts
                if date == end_ftime:
                    precip_fcst[model] = var_list[-1]
            # Try to fix max min issues
            if date >= start_ptime and date <= end_ptime:
                if max_val == None and min_val == None:
                    max_val = var_list[-1]
                    min_val = var_list[-1]
                if var_list[-1] > max_val:
                    max_val = var_list[-1]
                if var_list[-1] < min_val:
                    min_val = var_list[-1]
                 
        # Now plot
        if model.endswith('mos'):
            pylab.plot(keysorted,var_list,color=color_dict[model],linestyle='dashed',linewidth=2)
        else:
            pylab.plot(keysorted,var_list,color=color_dict[model],linewidth=2)

    # Special work here to use the GFS LAMP output, which has a different time
    # Series and needs little reformatting
    print "Processing LAMP guidance"
    try:
        lamp_fcst = lamp_parser(siteid.upper())
        keysorted = lamp_fcst['FDATE']
        if var == 'temp':
            var_list = lamp_fcst['TMP']
        elif var == 'wspd':
            var_list = lamp_fcst['WSP']
            # Also plot the gust as a dashed line
            gust_list = lamp_fcst['WGS']
            pylab.plot(keysorted,gust_list,color=color_dict['GFS LAMP'],linewidth=2,linestyle='dashed')
        elif var == 'precip':
            var_list = [np.nan for n in range(len(keysorted))]
        # Sort out what to call this
        lamp_start_hour = (keysorted[0] - timedelta(hours=1)).hour
        lamp_name = '%dZ GFS LAMP' % lamp_start_hour

        # Now plot this
        pylab.plot(keysorted,var_list,color=color_dict['GFS LAMP'],linewidth=2)
    except:
        lamp_name = 'GFS LAMP'

    # do some quick archiving of precip forecasts here
    if os.path.exists('%s/site_data/%s_precip_fcsts.pickle' % (maindir,siteid.upper())):
        out_precip = pickle.load(open('%s/site_data/%s_precip_fcsts.pickle' % (maindir,siteid.upper()),'r'))
    else:
        out_precip = {}
    pcp_fcst_date = datetime(start_ftime.year,start_ftime.month,start_ftime.day)
    for modnm in precip_fcst.keys():
        if model_ids[modnm] not in out_precip.keys():
            out_precip[model_ids[modnm]] = {}
        out_precip[model_ids[modnm]][pcp_fcst_date] = precip_fcst[modnm]
    pickle.dump(out_precip,open('%s/site_data/%s_precip_fcsts.pickle' % (maindir,siteid.upper()),'w'))

    # And the daily wind forecast
    # No overwriting here--new each time
    if var == 'wspd':
        pickle.dump(wind_fcsts,open('%s/site_data/%s_wind_fcsts.pickle' % (maindir,siteid.upper()),'w'))


    # Now plot the observations
    keysorted = oblist.keys()
    keysorted.sort()
    var_list = []
    for date in keysorted:
        if var == 'temp':
            var_list.append(C_to_F(oblist[date].tmpc))
        elif var == 'wspd':
            var_list.append(oblist[date].sknt)
        elif var == 'precip':
            if date <= start_ftime:
                var_list.append(0.0)
            else:
                var_list.append(var_list[-1]+mm_to_in(oblist[date].p01m))
    pylab.plot(keysorted,var_list,'ko:',linewidth=2,markersize=5) 
    ax = pylab.gca()


    # For temperature, add a text box with the MOS max and min
    if var == 'temp':
        txtb = ''
        for model in hilo_model.keys():
            hi,lo = hilo_model[model][start_ftime.date()][0:2]
            #try:
            txtb = txtb + '%s  H: %d  L: %d\n' % (model.upper(),hi,lo)
            
            #except:
                
        txtb = txtb[:-1]
        pylab.text(0.45,0.9,txtb,fontsize=14,transform = ax.transAxes,bbox=dict(facecolor='gray',alpha=0.1))

    unit_dict = {'temp':'F',
                 'wspd':'kts',
                 'precip':'in.'}

    # Some things to make the figure look nice
    pylab.title('Forecast %s at %s' % (var,siteid.upper()))
    pylab.xlabel('Valid Time')
    pylab.ylabel('%s [%s]' % (var.upper(),unit_dict[var]))
    legend_labels = [model_ids[mod] for mod in modlist]
    legend_labels = legend_labels + moslist + [lamp_name]
    if var == 'precip':
        # Put the legend at the top
        leg = pylab.legend(legend_labels,loc=9,ncol=6,mode='expand')
    elif var == 'wspd':
        # Put the legend in the best spot
        leg = pylab.legend(legend_labels,loc=0,ncol=6,mode='expand')
    else:
        # Put the legend at the bottom
        leg = pylab.legend(legend_labels,loc=8,ncol=6,mode='expand')
    # Make the legend transparent
    leg.get_frame().set_alpha(0.5)
    ltexts = leg.get_texts()
    pylab.setp(ltexts,fontsize='small')

    # Grey out non-forecast period
    if var == 'wspd':
        pylab.axvspan(matplotlib.dates.date2num(start_ptime),matplotlib.dates.date2num(start_ftime_wind),facecolor='0.8',alpha=0.40)
        pylab.axvspan(matplotlib.dates.date2num(end_ftime_wind),matplotlib.dates.date2num(end_ptime),facecolor='0.8',alpha=0.40)
    else:
        pylab.axvspan(matplotlib.dates.date2num(start_ptime),matplotlib.dates.date2num(start_ftime),facecolor='0.8',alpha=0.40)
        pylab.axvspan(matplotlib.dates.date2num(end_ftime),matplotlib.dates.date2num(end_ptime),facecolor='0.8',alpha=0.40)


    # Reformat the date axis
    ax.set_xlim([matplotlib.dates.date2num(start_ptime),matplotlib.dates.date2num(end_ptime)])
    #min,max = ax.get_ylim()
    if var == 'temp':
        ax.set_ylim((min_val-3,max_val+3))
        min,max = ax.get_ylim()
        ax.set_yticks(range(int(min),int(max)))
    elif var == 'wspd':
        ax.set_ylim((0,max_val+3))
        min,max = ax.get_ylim()
        ax.set_yticks(range(0,int(max)))

    ax.xaxis.set_major_locator(matplotlib.dates.HourLocator(byhour=[0,3,6,9,12,15,18,21]))
    ax.xaxis.set_major_formatter(matplotlib.dates.DateFormatter('%HZ'))
    ax.xaxis.set_minor_locator(matplotlib.dates.DayLocator())
    ax.xaxis.set_minor_formatter(matplotlib.dates.DateFormatter('%h %d'))
    pylab.grid()
    if export_flag:
        pylab.savefig('%s_%s.png' % (siteid.upper(),var),bbox_inches='tight')
        os.system('mv %s_%s.png %s/%s_%s.png' % (siteid.upper(),var,outdir,siteid.upper(),var))
    else:
        pylab.show()
        
# This will contain the high and low temperatures from MOS
hilo_model = {}

# Add the latest USL forecast
try:
    from scrub_USL import scrub_ALL_USL
    uslmod1,usl_hilo1,uslmod2,uslhilo2 = scrub_ALL_USL(siteid.upper()) 
    hilo_model['usl_mos'] = usl_hilo1
    if len(uslmod2.keys()) > 0:
        models['usl_p1'] = uslmod1
        models['usl'] = uslmod2
    else:
        models['usl'] = uslmod1
except:
    print "No USL data found."

modlist = models.keys()
model_init = {}

# Get Mos dictionary
moslist = ['gfs_mos','nam_mos']
final_moslist = []
for modtyp in moslist:
    try:
        models[modtyp],hilo_model[modtyp] = mos_parser(siteid,modtyp[0:3])
        final_moslist.append(modtyp)
    except:
        print "No MOS available locally for ", modtyp
moslist = final_moslist
print "MOS available:", moslist
for mos in moslist:
    print mos,hilo_model[mos]

# Get observations
oblist,xnlist = obs_parser(siteid)

moslist.sort()
modlist.sort()

# More refined figuring of the models

plot_wind_barbs(modlist,moslist)
plot_models('temp',modlist,oblist,moslist)
plot_models('wspd',modlist,oblist,moslist)
plot_models('precip',modlist,oblist,moslist)

